import os
import sys
cur_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(cur_dir)
sys.path.insert(0, os.path.abspath(os.path.join(cur_dir, '..')))

import paddle
from ppdet.core.workspace import load_config, create, merge_config
from ppdet.utils.cli import ArgsParser, merge_args
from ppdet.engine import Trainer
from ppdet.slim import build_slim_model

def parse_args():
    parser = ArgsParser()
    parser.add_argument("-p", "--pretrain_weights", default=None, help="Path to resume weights")
    parser.add_argument("-r", "--resume", default=None, help="Path to resume weights")
    parser.add_argument("--slim_config", default=None, type=str, help="Configuration file of slim method.")
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    try:
        cfg = load_config(args.config)
    except AttributeError:
        raise ValueError("Please provide a configuration file argument (--config or -c).")
    merge_args(cfg, args)
    merge_config(args.opt)
    if not os.path.exists(cfg.weights):
        os.makedirs(cfg.weights, exist_ok=True)
    
    # disable npu in config by default
    if 'use_npu' not in cfg:
        cfg.use_npu = False

    # disable xpu in config by default
    if 'use_xpu' not in cfg:
        cfg.use_xpu = False

    if 'use_gpu' not in cfg:
        cfg.use_gpu = False

    # disable mlu in config by default
    if 'use_mlu' not in cfg:
        cfg.use_mlu = False

    if cfg.use_gpu:
        place = paddle.set_device('gpu')
    elif cfg.use_npu:
        place = paddle.set_device('npu')
    elif cfg.use_xpu:
        place = paddle.set_device('xpu')
    elif cfg.use_mlu:
        place = paddle.set_device('mlu')
    else:
        place = paddle.set_device('cpu')    
        
    if args.slim_config:
        cfg = build_slim_model(cfg, args.slim_config)
    
    trainer = Trainer(cfg, mode='train')
    # 加载权重
    if cfg.resume is not None:
        trainer.resume_weights(cfg.resume)
    elif 'pretrain_weights' in cfg and cfg.pretrain_weights:
        trainer.load_weights(cfg.pretrain_weights)
    trainer.train()

if __name__ == '__main__':
    main()